'''
Author: Zhuangyu
Date: 2022-09-20 12:44:44
LastEditTime: 2023-03-22 21:14:12
'''
from Parameters import *
import van_Genuchten as vg
import numpy as np
from casadi import *
# from casadi import *

#The module contains the 1D Richards Model which is approximated 
#using the two point central difference approach.
#This is for a single soil type case.
#Will have to be modified when different soil types are being considered.
#The 1D model is expressed as c(h)dh/dt = d/dz(K(h)*d/dz(h+z)).

p=Loam()
Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()

def RichardsModel_1D(x,u,w,Kc,ET0):
    
    #x is the state vector
    #u is the irrigation amount (the input)
    #w is the process noise
    #Kc is the crop coefficient
    #ET0 is the reference evapotranspiration
    #p represents the soil parameters
    
    # dhdt=SX.zeros(Nz)
    dhdt=SX.zeros(Nz)
    for i in range(0,Nz):

        CurrentState=x[i]         
                
        if i == 0: #Bottom boundary, Free drainage boundary condition(BC)
            BCz_L=CurrentState
            BCz_U=x[i+1]
         
        elif i == Nz-1: #Top boundary
            BCz_L = x[i - 1]    
         
        else: #Internal nodes
            BCz_L = x[i - 1]
            BCz_U = x[i + 1]

        #Computing the unsaturated hydraulic conductivity at the centre of the compartments
        KzL=0.5*(vg.KFun(CurrentState,p)+vg.KFun(BCz_L,p))
        KzU=0.5*(vg.KFun(CurrentState,p)+vg.KFun(BCz_U,p))
        
        #Computing the pressure head gradient
        DHzL=(CurrentState-BCz_L)/dz 
        DHzU=(BCz_U-CurrentState)/dz

        if i == Nz-1: #Surface Node [Nuemann BC which incorporates the irrigation amount, u]
            Term1 = (1.0/dz)*( -u - KzL*DHzL)
            Term2 =(1.0/dz)*( -1*KzL)
        
        else: #Nodes below the surface
            Term1 = (1.0/dz) * (KzU*DHzU - KzL*DHzL)
            Term2 = (1.0/dz) * (KzU - KzL)

        Term3 = (Kc*ET0)/Hz #The Sink term [Root water extraction rate]
        Term4 = Term1 + Term2 - Term3
        Term5 = Term4/vg.CFun(CurrentState,p)

        dhdt[i]=Term5
    
    DHDT=dhdt + w
    return DHDT

def ODE_discrete(x,u,w,Kc,ET0,DeltaT):
    return x + DeltaT*RichardsModel_1D(x,u,w,Kc,ET0)

def getOutputs1D_allnodes(x):
    #Returns the measurements for all the states
    y = vg.thetaFun(x,p)
    CMatrix = cmatrix_single(Nz)
    # y_pred = mtimes(CMatrix,y)
    y_pred = mtimes(CMatrix,y)
    return y_pred

def getOutputs1D(x):
    #Returns measurements for some selected states based on the C-matrix
    y = vg.thetaFun(x,p)
    # C=np.hstack((np.zeros((N_obs,Nz-N_obs)),np.eye(N_obs)))
    C1=np.eye(Nz)
    C=np.delete(C1,(1,2,3,4,5,6,8,9,10,11,12,13,14),axis=0)
    # C=C1
    # y_pred = mtimes(C,y)
    y_pred = mtimes(C,y)    
    return y_pred
